import asyncio
import json
import time

from datasets import load_dataset

from lagent.agents import AsyncMathCoder
from lagent.agents.aggregator import InternLMToolAggregator
from lagent.llms import AsyncGPTAPI
from lagent.prompts.parsers import ToolParser

# 创建并设置异步事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

# 数学解题的提示模板
interpreter_prompt = (
    # 要求使用 Python 代码逐步解决数学问题
    'Below is a math problem. Please solve it step by step with the assistance of Python programming. Consider using Sympy or Numpy library '
    # 建议使用 Sympy 处理符号计算
    'to facilitate your derivation, calculation and equation solving. Utilize the "pi" symbol and "Rational" from Sympy '
    # 要求使用分数和根号形式，而不是小数
    'for $$\pi$$ and fractions, and simplify all fractions and square roots without converting them to decimal values. '
    # 定义代码块的标记格式
    'Please encapsulate each generated Jupyter Python code block with tags "{begin}" and "{end}". Conclude the '
    # 要求使用 LaTeX 格式输出最终答案
    r'final answer when observations are sufficient and encapsulate the numerical result with LaTeX syntax "\boxed{{}}" '
    # 使用 [END] 标记结束
    'without any unit, and end your conclusion with the special token "[END]" to denote the completion of your response. '
    'Keep the following points in mind:\n'
    # 要求交替使用自然语言和代码
    '- You must alternately use human and programming languages in the chain of thought;\n'
    # 限制推理步骤不超过三步
    '- The number of your reasoning steps should not exceed **three**, which means you may merge some intermediate steps when the original answer is tedious.'
)

# 配置 GPT API
async_llm = dict(
    type=AsyncGPTAPI,
    model='gpt-4o-2024-05-13',  # 使用 GPT-4 模型
    retry=50,  # 最大重试次数
    key='',  # API 密钥
    max_new_tokens=2048,  # 最大生成长度
    stop_words=['</python'],  # 停止词
    proxies=dict(),  # 代理设置
)

# 创建异步数学解题代理
async_agent = AsyncMathCoder(
    llm=async_llm,
    output_format=ToolParser(
        tool_type='interpreter',  # 使用解释器工具
        template=interpreter_prompt,  # 使用上面定义的提示模板
        begin='<python>\n',  # 代码块开始标记
        end='\n</python>'  # 代码块结束标记
    ),
    aggregator=InternLMToolAggregator(
        environment_role='system',  # 环境角色设置
        environment_begin='<output>\n',  # 输出开始标记
        environment_end='\n</output>'  # 输出结束标记
    ),
    finish_condition=lambda m: '[END]' in m.content,  # 定义完成条件
)

# 加载数学问题数据集
ds = load_dataset('lighteval/MATH', split='train')
problems = [item['problem'] for item in ds.select(range(30))]  # 选取前30个问题

# 开始批量解题
tic = time.time()  # 记录开始时间
# 为每个问题创建一个异步任务
coros = [async_agent(q, session_id=i) for i, q in enumerate(problems)]
# 并发执行所有任务
res = loop.run_until_complete(asyncio.gather(*coros))
print(time.time() - tic)  # 打印总耗时

# 保存解题步骤到文件
with open('tmp_6.json', 'w') as f:
    json.dump([async_agent.get_steps(i) for i in range(len(problems))],
              f,
              ensure_ascii=False,
              indent=4)
